import torch

from utils.randomUtils import randomSelectFromTensor, weightedRandomSelect


class ResultRecorder:

    def __init__(self, maxNum, valSetSize, threshold):
        # maxNum，可以容纳的最大记录条数
        self.maxNum = maxNum
        self.size = 0
        self.valSetSize = valSetSize
        # 记录正确结果的Tensor，1表示预测正确，0表示预测错误。
        self.corrTensor = torch.zeros((maxNum, valSetSize), dtype=torch.int8)
        # 记录正确结果的confidence，范围（0，,1）
        self.softMaxResult = torch.zeros(maxNum, valSetSize, dtype=torch.float64)
        # threshold稳定数据集的判定标准(0~1)
        self.threshold = threshold
        # 全量记录
        self.fullIndex = torch.arange(0, self.valSetSize)

    def getHardIndicesFromCorrTensor(self):
        """
        根据历史结果生成筛选结果
        :return: 返回待检测的数据集
        """
        if self.size != self.maxNum:
            # 返回全部的序列
            return [i for i in range(self.valSetSize)]
        success_rate = torch.sum(self.corrTensor, dim=0) / self.maxNum
        # 使用.nonzero()方法找到True值的位置,并拉平成列表
        to_val_indices = (success_rate < self.threshold).nonzero(as_tuple=False).view(-1)
        return to_val_indices

    def getEasyIndicesFromCorrTensor(self):
        """
        根据历史结果生成筛选结果
        :return: 返回待检测的数据集
        """
        if self.size != self.maxNum:
            # 返回全部的序列
            return [i for i in range(self.valSetSize)]
        success_rate = torch.sum(self.corrTensor, dim=0) / self.maxNum
        # 使用.nonzero()方法找到True值的位置,并拉平成列表
        to_val_indices = (success_rate >= self.threshold).nonzero(as_tuple=False).view(-1)
        return to_val_indices

    def addNewCorrRecord(self, corrRecord):
        """
        添加一条新的记录到记录器中
        :param corrRecord:是一个二维tensor[]，(1,n)
        :return:
        """
        # 确保 CorrRecord 是 PyTorch 张量，并且形状正确
        if self.size < self.maxNum:
            self.size += 1
        # 将一维张量扩展为二维张量
        corrRecord = corrRecord.unsqueeze(0)  # 形状变为 (1, -1)
        self.corrTensor = torch.cat((self.corrTensor[1:], corrRecord), dim=0)

    def addNewConfidenceRecord(self, confidenceRecord):
        # 将一维张量扩展为二维张量
        confidenceRecord = confidenceRecord.unsqueeze(0)  # 形状变为 (1, -1)
        self.softMaxResult = torch.cat((self.softMaxResult[1:], confidenceRecord), dim=0)

    def getIndicesWithRandomChoose(self, sample_ratio=None):
        """
        根据maxNum从忽略的样本集中随机抽样一定比例的数据加入到选中的样本集中
        抽取的比例默认为0.1
        """
        if sample_ratio == 1.0 or self.size != self.maxNum:
            return self.fullIndex
        if sample_ratio is None:
            sample_ratio = 0.1
        hard_indices = self.getHardIndicesFromCorrTensor()
        easy_indices = self.getEasyIndicesFromCorrTensor()
        print(f"hard_indices={len(hard_indices)},easy_indices={len(easy_indices)}")
        sample_num = int(len(easy_indices) * sample_ratio)
        recover_indices = randomSelectFromTensor(easy_indices, sample_num)
        combined = torch.cat((hard_indices, recover_indices))
        return combined

    def getIndicesWithWeightRandomChoose(self, sample_ratio=None):
        if sample_ratio == 1.0 or self.size != self.maxNum:
            return self.fullIndex
        if sample_ratio is None:
            sample_ratio = 0.1
        hard_indices = self.getHardIndicesFromCorrTensor()
        easy_indices = self.getEasyIndicesFromCorrTensor()

        sample_num = int(len(easy_indices) * sample_ratio)+1
        if len(easy_indices) == 0 or sample_num == 0:
            return hard_indices
        mean_confidences = torch.mean(self.softMaxResult, dim=0, keepdim=False)
        easy_confidences = mean_confidences[easy_indices]
        easy_weights = 1 / easy_confidences
        print(f'easy_weights={easy_weights}')
        weight_indices = weightedRandomSelect(easy_weights, sample_num)
        recover_indices = easy_indices[weight_indices]
        combined = torch.cat((hard_indices, recover_indices))
        return combined

    def getLatestConfidenceRecord(self):
        return self.softMaxResult[-1].clone()
